from torch.utils.data import Dataset, DataLoader
import math


class My_DataSet(Dataset):
    """需要实现getitem和len魔法方法"""

    def __init__(self):
        # 获取数据
        # readlines方法是将文件中的数据按行都读出来存入列表中
        self.lines = open(r"F:\virtual_environment\data\SMSSpamCollection", mode="r", encoding="UTF-8").readlines()

    def __getitem__(self, index):  # 该方法可以使用对象[index]方法取值
        #         # strip是删除字符串开头和末尾指定的字符串，当不指定时是空格
        #         return self.lines[index].strip()
        label = self.lines[index][:4].strip()
        context = self.lines[index][4:].strip()
        return label, context

    def __len__(self):
        return len(self.lines)


if __name__ == '__main__':

    dataset = My_DataSet()
    # print(dataset[1])
    # 参数依次是：继承自DataSet类的对象；每个batch中的数据条数；是否打乱顺序，默认是False；加载数据的线程数量
    data_loader = DataLoader(dataset=dataset, batch_size=2, shuffle=True, num_workers=2)
    print(len(dataset))
    print(len(data_loader))  # len(dataset)/batch_size
    print(math.ceil(len(dataset) / 2))
    for i in data_loader:
        print(i)



